# scripts/update_daily.py
#!/usr/bin/env python3
"""
更新日线数据脚本
"""
import os
import sys
import argparse
from pathlib import Path
import time

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

def update_daily_data():
    """更新日线数据"""
    parser = argparse.ArgumentParser(description='更新日线数据')
    parser.add_argument('--days', type=int, default=1, help='更新最近几天数据')
    parser.add_argument('--symbol', type=str, help='指定股票代码')
    parser.add_argument('--all', action='store_true', help='更新所有股票')
    parser.add_argument('--batch-size', type=int, default=100, help='批量处理数量')

    args = parser.parse_args()

    from stock_data_system.main import setup_logging, AppConfig, StockDataService
    import logging

    setup_logging(logging.INFO)
    logger = logging.getLogger(__name__)

    print("=" * 50)
    print("日线数据更新脚本")
    print("=" * 50)

    try:
        config = AppConfig()
        service = StockDataService(config.database, config.baostock)

        # 检查数据库连接
        if not service.db_manager.check_connection():
            print("❌ 数据库连接失败！")
            return 1

        print("✅ 数据库连接成功")

        start_time = time.time()

        if args.symbol:
            # 更新指定股票
            print(f"正在更新 {args.symbol} 的最近 {args.days} 天数据...")
            symbols = [args.symbol]
            count = service.update_daily_data(symbols, days=args.days)

            if count > 0:
                print(f"✅ 成功更新 {args.symbol} 的 {count} 条日线数据")
            else:
                print(f"⚠️  {args.symbol} 没有新数据")

        elif args.all:
            # 更新所有股票
            print(f"正在更新所有股票的最近 {args.days} 天数据...")
            count = service.update_all_stocks_daily(args.days)
            print(f"✅ 成功更新 {count} 条日线数据")

        else:
            print("请指定 --symbol 股票代码或 --all 更新所有股票")
            parser.print_help()
            return 1

        elapsed_time = time.time() - start_time
        print(f"⏱️  总耗时: {elapsed_time:.2f} 秒")

        return 0

    except KeyboardInterrupt:
        print("\n⚠️  用户中断操作")
        return 130
    except Exception as e:
        print(f"❌ 更新失败: {e}")
        return 1

if __name__ == "__main__":
    sys.exit(update_daily_data())